#include "PDTWhile.h"
#include "PDTCodeGraph.h"
#include "PDTZ3GraphVisitor.h"

void
PDTWhile::setCondition(PDTExpBool & condition)
{
  addChild(&condition);
  condition.addParent(this);
}

void
PDTWhile::setLoopStatement(PDTStatement & loopStatement)
{
  addChild(&loopStatement);
  loopStatement.addParent(this);
}

PDTExpBool*
PDTWhile::getCondition() const
{
  return dynamic_cast<PDTExpBool *> (children[0]);
}

PDTStatement *
PDTWhile::getLoopStatement() const
{
  return dynamic_cast<PDTStatement *> (children[1]);
}

PDTWhile::PDTWhile(PDTCodeGraph & CG, PDTExpBool & condition,
    PDTStatement & loopStatement) :
  PDTStatement(CG)
{
  addChild(&condition);
  condition.addParent(this);

  addChild(&loopStatement);
  loopStatement.addParent(this);

  CG.maxWhileDepth++;
  incrementChildrenDepth();
}

Z3_ast
PDTWhile::Z3NodeVisit (PDTZ3Visitor &ZV)
{
  if (Z3node != NULL)
    return Z3node;
#if 0		
  (define-sorts ((b (Array Int statement))))
  (define-sorts ((_startWhile (Array Int statement))))
  (define-sorts ((s1 (Array Int statement))))
  (define-sorts ((s2 (Array Int statement))))
  (define-sorts ((_endOfWhile (Array Int statement))))

  (declare_const _beforeWhile statement)
  (declare_const _finishWhile statement)
  (declare_const Infinity statement)
#endif
  /* declare arrays of statements to match the following setup
   * beforeW
   * while ( cond[] guard ) {
   *   startWhile []
   *   body []
   *   endWhile []
   * }
   * finishWhile
   */
  PDTNode * loopCond = getCondition();

  Z3_ast Z3beforeWhile, Z3finishWhile;
  Z3_ast Z3startWhile, Z3endOfWhile;
  Z3_ast Z3cond, Z3body;

  char * name = makeUniqueName("wGuard");
  Z3cond = ZV.makeArrayOfStatements(depth, name);

  name = makeUniqueName("wStart");
  Z3startWhile = ZV.makeArrayOfStatements(depth, name);

  name = makeUniqueName("wBody");
  Z3body = ZV.makeArrayOfStatements(depth, name);//TODO: ?? what is this?

  name = makeUniqueName("wEnd");
  Z3endOfWhile = ZV.makeArrayOfStatements(depth, name);

  // current?

  name = makeUniqueName("wBefore");
  Z3beforeWhile = ZV.makeStatement(name);

  name = makeUniqueName("wFinish");
  Z3finishWhile = ZV.makeStatement(name);

  /*
   * Set ZV.current to the statement before the while condition!
   * TODO: is this correct? didn't we just lose the actual statement before the while?
   */
  ZV.current = Z3beforeWhile;

  // forall i, Z3cond[i] = getcondition()->Z3visit()

  Z3_ast termsCond, bodyCond, Z3cond_i, Z3condNode;
  Z3_pattern patternsCond[1];
  Z3_sort sortsCond[1];
  Z3_symbol namesCond[1];

  termsCond = Z3_mk_bound(ZV.ctx, 0, ZV.Int);
  patternsCond[0] = Z3_mk_pattern(ZV.ctx, 1, &termsCond);
  sortsCond[0] = ZV.Int;

  /*
   * Build a variable to quantify statements
   */
  Z3_ast iCond; /* TODO: this is a variable iterating in the forall loop */
  name = makeUniqueName("wICond");
  namesCond[0] = Z3_mk_string_symbol(ZV.ctx, name);
  iCond = Z3_mk_const(ZV.ctx, namesCond[0], ZV.Int);

  /* Z3cond_i = Z3cond[iCond]*/
  Z3cond_i = Z3_mk_select(ZV.ctx, Z3cond, iCond);

  /* TODO: make sure the following is correct
   * The intention is that the conditional expression will use Z3cond_i as the statement in its assertions.
   * However all its assert claims need to happen within a forall
   * */
  ZV.current = Z3cond_i;
  Z3condNode = getCondition()->Z3visit(ZV);

  Z3_ast s2i_Apply, s2i_args[1];
  s2i_args[0] = Z3cond_i;
  s2i_Apply = Z3_mk_app(ZV.ctx, ZV.s2i[depth], 1, s2i_args);
  bodyCond = Z3_mk_eq(ZV.ctx, s2i_Apply, Z3condNode);

  Z3_ast zero, firstCond, constraints[8];

  zero = Z3_mk_int(ZV.ctx, 0, ZV.Int);
  firstCond = Z3_mk_select(ZV.ctx, Z3cond, zero);
  ZV.assertOrder(Z3beforeWhile, firstCond);

  constraints[0] = Z3_mk_forall(ZV.ctx, 0, 1, patternsCond, 1, sortsCond,
      namesCond, bodyCond);

  // order between condition and startWhile

  Z3_symbol namesOrder[1];
  Z3_ast iOrder2, Z3cond_i2, Z3startWhile_i2, E_cond_i2, E_startWhile_i2,
      bodyOrder2;
  Z3_ast argsCond_i2[1], argsStartWhile_i2[1];

  namesOrder[0] = Z3_mk_string_symbol(ZV.ctx, "iOrder2");
  iOrder2 = Z3_mk_const(ZV.ctx, namesOrder[0], ZV.Int);

  Z3cond_i2 = Z3_mk_select(ZV.ctx, Z3cond, iOrder2);
  Z3startWhile_i2 = Z3_mk_select(ZV.ctx, Z3startWhile, iOrder2);

  argsCond_i2[0] = Z3cond_i2;
  argsStartWhile_i2[0] = Z3startWhile_i2;

  E_cond_i2 = Z3_mk_app(ZV.ctx, ZV.E, 1, argsCond_i2);
  E_startWhile_i2 = Z3_mk_app(ZV.ctx, ZV.E, 1, argsStartWhile_i2);

  bodyOrder2 = Z3_mk_lt(ZV.ctx, E_cond_i2, E_startWhile_i2);

  Z3_ast terms2;
  Z3_pattern patterns2[1];
  Z3_sort sorts2[1];

  terms2 = Z3_mk_bound(ZV.ctx, 0, ZV.Statement);
  patterns2[0] = Z3_mk_pattern(ZV.ctx, 1, &terms2);
  sorts2[0] = ZV.Statement;

  constraints[1] = Z3_mk_forall(ZV.ctx, 0, 1, patterns2, 1, sorts2, namesOrder,
      bodyOrder2);

  // current = startWhile[] (array)

  // using the same Z3_ast node created above Z3startWhile_i2:
  Z3_ast bodyCurrent1;
  bodyCurrent1 = Z3_mk_eq(ZV.ctx, ZV.current, Z3startWhile_i2);/* TODO: make sure it's correct*/
  constraints[2] = Z3_mk_forall(ZV.ctx, 0, 1, patterns2, 1, sorts2, namesOrder,
      bodyCurrent1);

  // listOfStatements->Z3visit()

  PDTNode * loopBody = getLoopStatement();
  loopBody->Z3visit(ZV);

  // order between current and endWhile

  //also using same iterator as above called iOrder2
  Z3_ast bodyEndOfWhile, Z3endOfWhile_i2, argsEndOfWhile_i2[1],
      argsCurrent_i2[1];
  Z3_ast E_endOfWhile_i2, E_current_i2;
  Z3endOfWhile_i2 = Z3_mk_select(ZV.ctx, Z3endOfWhile, iOrder2);
  argsEndOfWhile_i2[0] = Z3endOfWhile_i2;
  argsCurrent_i2[0] = ZV.current;

  E_endOfWhile_i2 = Z3_mk_app(ZV.ctx, ZV.E, 1, argsEndOfWhile_i2);
  E_current_i2 = Z3_mk_app(ZV.ctx, ZV.E, 1, argsCurrent_i2);
  bodyEndOfWhile = Z3_mk_lt(ZV.ctx, ZV.current, Z3endOfWhile_i2);

  // order between endWhile, cond, finishWhile


  /*(assert (forall (i Int) (if (> i 0) (> (E (select b i)) (E (select _endOfWhile (- i 1))))))) */
  Z3_ast it1, condition1, result1, Z3cond_it1, Z3endOfWhile_it1_minus_1, minus;
  Z3_ast E_cond_it1, E_endOfWhile_it1_minus_1, implies;
  Z3_ast args3[1], args4[2], args5[1];
  Z3_symbol names3[1], one;

  names3[0] = Z3_mk_string_symbol(ZV.ctx, "it1");
  it1 = Z3_mk_const(ZV.ctx, names3[0], ZV.Int);

  Z3cond_it1 = Z3_mk_select(ZV.ctx, Z3cond, it1);
  args3[0] = Z3cond_it1;
  E_cond_it1 = Z3_mk_app(ZV.ctx, ZV.E, 1, args3);

  one = Z3_mk_int_symbol(ZV.ctx, 1);
  args4[0] = it1;
  args4[1] = Z3_mk_const(ZV.ctx, one, ZV.Int);
  minus = Z3_mk_sub(ZV.ctx, 2, args4);

  Z3endOfWhile_it1_minus_1 = Z3_mk_select(ZV.ctx, Z3endOfWhile, minus);
  args5[0] = Z3endOfWhile_it1_minus_1;
  E_endOfWhile_it1_minus_1 = Z3_mk_app(ZV.ctx, ZV.E, 1, args5);

  condition1 = Z3_mk_gt(ZV.ctx, it1, zero);
  result1 = Z3_mk_gt(ZV.ctx, E_cond_it1, E_endOfWhile_it1_minus_1);

  implies = Z3_mk_implies(ZV.ctx, condition1, result1);

  Z3_ast terms3;
  Z3_pattern patterns3[1];
  Z3_sort sorts3[1];

  terms3 = Z3_mk_bound(ZV.ctx, 0, ZV.Int);
  patterns3[0] = Z3_mk_pattern(ZV.ctx, 1, &terms3);
  sorts3[0] = ZV.Int;

  constraints[3] = Z3_mk_forall(ZV.ctx, 0, 1, patterns3, 1, sorts3, names3,
      implies);

  /*;;;; if !b[0] then E(finish) > E(b[0]) and E(finish)<E(infinity)
   (assert (> (E _finishWhile) (E (select b 0))))*/
  Z3_ast E_finishWhile, E_firstCond, args6[1], args7[1];

  args6[0] = firstCond;
  args7[0] = Z3finishWhile;
  E_firstCond = Z3_mk_app(ZV.ctx, ZV.E, 1, args6);
  E_finishWhile = Z3_mk_app(ZV.ctx, ZV.E, 1, args7);

  constraints[4] = Z3_mk_gt(ZV.ctx, E_finishWhile, E_firstCond);

  /*(assert (=> (not (= (s2i (select b 0)) 1)) )   (< (E _finishWhile) (E Infinity)))*/

  Z3_ast s2i_firstCond, E_Infinity, condition2, result2, args8[1];

  args8[0] = ZV.Infinity;
  s2i_firstCond = Z3_mk_app(ZV.ctx, ZV.s2i[depth], 1, args6);
  E_Infinity = Z3_mk_app(ZV.ctx, ZV.E, 1, args8);
  condition2 = Z3_mk_eq(ZV.ctx, s2i_firstCond, ZV.False);
  result2 = Z3_mk_lt(ZV.ctx, E_finishWhile, E_Infinity);

  constraints[5] = Z3_mk_implies(ZV.ctx, condition2, result2);

  /*;;;; forall i if b[i] then E(finish)>E(b[i])
   (assert (forall (i Int) (=> (= (s2i (select b i)) 1)) (> (E _finishWhile) (E select b i))) ))*/

  Z3_ast it2, cond_it2, s2i_cond_it2, E_cond_it2, args9[1], implies3;
  Z3_ast condition3, result3;
  Z3_symbol names4[1];

  names4[0] = Z3_mk_string_symbol(ZV.ctx, "it2");
  it2 = Z3_mk_const(ZV.ctx, names4[0], ZV.Int);

  cond_it2 = Z3_mk_select(ZV.ctx, Z3cond, it2);
  args9[0] = cond_it2;
  E_cond_it2 = Z3_mk_app(ZV.ctx, ZV.E, 1, args9);
  s2i_cond_it2 = Z3_mk_app(ZV.ctx, ZV.s2i[depth], 1, args9);

  condition3 = Z3_mk_eq(ZV.ctx, s2i_cond_it2, ZV.True);
  result3 = Z3_mk_gt(ZV.ctx, E_finishWhile, E_cond_it2);
  implies3 = Z3_mk_implies(ZV.ctx, condition3, result3);

  Z3_ast terms4;
  Z3_pattern patterns4[1];
  Z3_sort sorts4[1];

  terms4 = Z3_mk_bound(ZV.ctx, 0, ZV.Int);
  patterns4[0] = Z3_mk_pattern(ZV.ctx, 1, &terms4);
  sorts4[0] = ZV.Int;

  constraints[6] = Z3_mk_forall(ZV.ctx, 0, 1, patterns4, 1, sorts4, names4,
      implies3);

  /*;;;; forall i, if (i>0 and !b[i] and b[i-1]), E(finish)>E(b[i]) and E(finish)<E(infinity)
   (assert (forall (i Int) (=> (and (> i 0) (= (s2i (select b i)) 0) (= (s2i (select b (- i 1))) 1))
   (and (> (E _finishWhile) (E (select b i))) (< (E _finishWhile) (E Infinity))))))*/

  Z3_ast it3, s2i_cond_it3, cond_it3_minus_1, s2i_cond_it3_minus_1, cond_it3;
  Z3_ast args10[1], args11[2], E_cond_it3, and1, and2, implies4, minus2;
  Z3_ast args12[1], conditions[3], results[2];
  Z3_symbol names5[1];

  names5[0] = Z3_mk_string_symbol(ZV.ctx, "it3");
  it3 = Z3_mk_const(ZV.ctx, names5[0], ZV.Int);

  cond_it3 = Z3_mk_select(ZV.ctx, Z3cond, it3);
  args10[0] = cond_it3;
  E_cond_it3 = Z3_mk_app(ZV.ctx, ZV.E, 1, args10);
  s2i_cond_it3 = Z3_mk_app(ZV.ctx, ZV.s2i[depth], 1, args10);

  args11[0] = it3;
  args11[1] = Z3_mk_const(ZV.ctx, one, ZV.Int);
  minus2 = Z3_mk_sub(ZV.ctx, 2, args11);

  cond_it3_minus_1 = Z3_mk_select(ZV.ctx, Z3cond, minus2);
  args12[0] = cond_it3_minus_1;
  s2i_cond_it3_minus_1 = Z3_mk_app(ZV.ctx, ZV.s2i[depth], 1, args12);

  conditions[0] = Z3_mk_gt(ZV.ctx, it3, zero);
  conditions[1] = Z3_mk_eq(ZV.ctx, s2i_cond_it3, ZV.False);
  conditions[2] = Z3_mk_eq(ZV.ctx, s2i_cond_it3_minus_1, ZV.True);

  results[0] = Z3_mk_gt(ZV.ctx, E_finishWhile, E_cond_it3);
  results[1] = Z3_mk_lt(ZV.ctx, E_finishWhile, E_Infinity);

  and1 = Z3_mk_and(ZV.ctx, 3, conditions);
  and2 = Z3_mk_and(ZV.ctx, 2, results);
  implies4 = Z3_mk_implies(ZV.ctx, and1, and2);

  Z3_ast terms5;
  Z3_pattern patterns5[1];
  Z3_sort sorts5[1];

  terms5 = Z3_mk_bound(ZV.ctx, 0, ZV.Int);
  patterns5[0] = Z3_mk_pattern(ZV.ctx, 1, &terms5);
  sorts5[0] = ZV.Int;

  constraints[7] = Z3_mk_forall(ZV.ctx, 0, 1, patterns5, 1, sorts5, names5,
      implies4);

  // current = finishWhile

  ZV.current = Z3finishWhile;
  Z3node = ZV.current;
  return Z3node;

# if 0
  (assert (> (E (select b 0)) (E _beforeWhile)))

  (assert (forall (i Int) (> (E (select _startWhile i)) (if (= (s2i (select b i)) 1) (E (select b i)) (E Infinity)))))
  (assert (forall (i Int) (> (E (select s1 i)) (E (select _startWhile i)))))
  (assert (forall (i Int) (> (E (select s2 i)) (E (select s1 i)))))
  (assert (forall (i Int) (> (E (select _endOfWhile i)) (E (select s2 i)))))
  (assert (forall (i Int) (if (> i 0) (> (E (select b i)) (E (select _endOfWhile (- i 1)))))))
#endif

#if 0
  ;;;; if !b[0] then E(finish) > E(b[0]) and E(finish)<E(infinity)

  (assert (> (E _finishWhile) (E (select b 0))))
  (assert (=> (not (= (s2i (select b 0)) 0)) ) (< (E _finishWhile) (E Infinity)))

  ;;;; forall i if b[i] then E(finish)>E(b[i])

  (assert (forall (i Int) (=> (= (s2i (select b i)) 1)) (> (E _finishWhile) (E select b i))) ))

  ;;;; forall i, if (i>0 and !b[i] and b[i-1]), E(finish)>E(b[i]) and E(finish)<E(infinity)

  (assert (forall (i Int) (=> (and (> i 0) (= (s2i (select b i)) 0) (= (s2i (select b (- i 1))) 1)) (and (> (E _finishWhile) (E (select b i))) (< (E _finishWhile) (E Infinity))))))
#endif


}
